# -*- coding: utf-8 -*-

# @File  : db.py
# @Author: Lomo
# @Site  : lomo.space
# @Date  : 2020-02-10
# @Desc  : DB MODEL


import pymysql

from config.db import DB_CONFIG


class DB(object):
    def __init__(self, _type='LOCAL'):
        self.host = DB_CONFIG.get(_type).get('host')
        self.port = DB_CONFIG.get(_type).get('port')
        self.user = DB_CONFIG.get(_type).get('user')
        self.password = DB_CONFIG.get(_type).get('password')
        self.database = DB_CONFIG.get(_type).get('database')

        self.conn = pymysql.connect(host=self.host, port=self.port, user=self.user,
                                    password=self.password, db=self.database)
        self.cur = self.conn.cursor(cursor=pymysql.cursors.DictCursor)

    def __enter__(self):
        return self.cur

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.conn.commit()
        self.cur.close()
        self.conn.close()

